import org.xbill.DNS.*;
import org.xbill.DNS.Record;

import java.io.IOException;
import java.net.*;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;


public class Handler implements Runnable{
        DatagramPacket request;
        DatagramSocket socket;
        Cache cache;
        Log log;
        String remoteDnsServer;
        boolean useCache;

        public Handler(DatagramPacket request, DatagramSocket socket, Cache cache, Log log, String remoteDnsServer, boolean useCache) {
            this.request = request;
            this.socket = socket;
            this.log = log;
            this.remoteDnsServer = remoteDnsServer;
            this.useCache = useCache;
            if (useCache) {
                this.cache = cache;
            }
        }

        @Override
        public void run() {
            InetAddress srcIp = request.getAddress();
            int srcPort = request.getPort();
            Message messageIn;
            try {
                messageIn = new Message(request.getData());
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            Record question = messageIn.getQuestion();
            String domain = question.getName().toString();
            boolean valid = true, useV6 = false, nop = false;

            int type = question.getType();
            switch (type) {
                // inverse dns
                case 12 -> {
                    nop = true;
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "inverse dns, simple relay");
                }
                // ipv4 question
                case 1 -> {
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "ipv4 question for domain: " + domain);
                }
                // ipv6 question
                case 28 -> {
                    useV6 = true;
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "ipv6 question for domain: " + domain);
                }
                default -> {
                    nop = true;
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "other type, simple relay");
                }
            }

            InetAddress ansIp = null;
            DatagramPacket response = null;

            if (useCache) {
                ansIp = cache.getIpFromCache(domain + (useV6 ? "-v6" : ""));
            }
            if (!nop && (ansIp != null)) {
                log.addLog("[" + Thread.currentThread().getName() + "] " + "found in cache");
            }
            else
            {
                if (!nop && useCache) {
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "not in cache");
                }
                DatagramSocket relaySocket;
                try {
                    relaySocket = new DatagramSocket();
                } catch (SocketException e) {
                    throw new RuntimeException(e);
                }
                byte[] relayBuf = messageIn.toWire();
                InetAddress dnsSeverIp;
                try {
                    dnsSeverIp = InetAddress.getByName(remoteDnsServer);
                } catch (UnknownHostException e) {
                    throw new RuntimeException(e);
                }
                DatagramPacket relayRequest = new DatagramPacket(relayBuf, relayBuf.length, dnsSeverIp, 53);
                byte[] buf = new byte[1024];
                DatagramPacket relayResponse = new DatagramPacket(buf, buf.length);

                try {
                    relaySocket.send(relayRequest);
                    relaySocket.receive(relayResponse);
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
                relaySocket.close();

                if (nop) {
                    response = relayResponse;
                } else {
                    Message messageResponse;
                    try {
                        messageResponse = new Message(relayResponse.getData());
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                    List<Record> records = messageResponse.getSection(Section.ANSWER);
                    ArrayList<InetAddress> ips = new ArrayList<>();
                    for (Record record : records) {
                        if (!useV6 && record instanceof ARecord) {
                            // ipv4 records
                            ARecord aRecord = (ARecord) record;
                            try {
                                InetAddress ip = InetAddress.getByAddress(aRecord.getAddress().getAddress());
                                ips.add(ip);
                            } catch (UnknownHostException e) {
                                throw new RuntimeException(e);
                            }
                        } else if (useV6 && record instanceof AAAARecord) {
                            // ipv6 records
                            AAAARecord aaaaRecord = (AAAARecord) record;
                            try {
                                InetAddress ip = InetAddress.getByAddress(aaaaRecord.getAddress().getAddress());
                                ips.add(ip);
                            } catch (UnknownHostException e) {
                                throw new RuntimeException(e);
                            }
                        }
                    }
                    if (ips.size() == 0) {
                        log.addLog("[" + Thread.currentThread().getName() + "] " + "no ipv" + (useV6 ? 6 : 4) + " result found from remote dns");
                        valid = false;
                    } else {
                        log.addLog("[" + Thread.currentThread().getName() + "] " + "in total " + ips.size() + " result(s)");
                        ansIp = ips.get(new Random().nextInt(ips.size()));
                        if (useCache) {
                            boolean finalUseV6 = useV6;
                            String parentName = Thread.currentThread().getName();
                            Thread update = new Thread(new Runnable() {
                                @Override
                                public void run() {
                                    synchronized (cache.getCacheLock()) {
                                        if (cache.getIpFromCache(domain + (finalUseV6 ? "-v6" : "")) == null) {
                                            cache.addCacheToFile(domain + (finalUseV6 ? "-v6" : ""), ips);
                                            log.addLog("[" + parentName + "-child] " + "added to cache file and reloaded cache");
                                        }
                                    }

                                }
                            });
                            update.start();
                        }
                    }
                }
            }
            if (!nop) {
                Message messageOut = messageIn.clone();
                if (!valid || ansIp.toString().substring(1).equals("0.0.0.0")
                        || ansIp.toString().substring(1).equals("::")
                        || ansIp.toString().substring(1).equals("0:0:0:0:0:0:0:0")) {
                    messageOut.getHeader().setRcode(3);
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "answer: non-existent domain");
                } else {
                    Record answer;
                    // ipv4 answer
                    if (!useV6) {
                        answer = new ARecord(question.getName(), question.getDClass(), 64, ansIp);
                    }
                        // ipv6 answer
                    else {
                        answer = new AAAARecord(question.getName(), question.getDClass(), 64, ansIp);
                    }
                    messageOut.addRecord(answer, Section.ANSWER);
                    log.addLog("[" + Thread.currentThread().getName() + "] " + "answer ip: " + ansIp.toString().substring(1));
                }
                byte[] buf = messageOut.toWire();
                response = new DatagramPacket(buf, buf.length);
            }
            response.setAddress(srcIp);
            response.setPort(srcPort);
            try {
                socket.send(response);
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
            log.writeLog();
        }
    }


